import cv2
import numpy as np


class RectangleBboxGenerator:
    def __init__(self, pixel, threshold=0.6,
                 min_pixel_size=10,
                 is_corrected=True,
                 enlarge_bbox_ratio=4,
                 enlarge_bbox_pixel=1):
        """

        :param pixel:
        :param threshold:
        :param min_pixel_size:
        """
        self.enlarge_bbox_pixel = enlarge_bbox_pixel
        self.enlarge_bbox_ratio = enlarge_bbox_ratio
        self.min_pixel_size = min_pixel_size
        self.pixel_threshold = threshold
        self.is_corrected = is_corrected
        self.text_area_pixel = pixel
        self.format_pixel = self._format_pixel()
        self.connected_domain = self.__get_connected_domains()

    def gen_bboxes(self):
        """
        interface
        :return:
        """
        domain_max_value = self.connected_domain.max()
        if domain_max_value == 0:
            return []
        bboxes = []
        for i in range(1, domain_max_value + 1):  # for each domain
            box = self._get_single_bbox(i)
            if box is not None:
                box = np.array(box) - 1
                hs = box[:, 1]
                ws = box[:, 0]
                box = np.array(
                    [[ws.min(), hs.min()], [ws.max(), hs.min()], [ws.max(), hs.max()], [ws.min(), hs.max()]])
                box = box * self.enlarge_bbox_ratio
                bboxes.append(box)
        return bboxes

    def __get_connected_domains(self):
        """
        combine and return connected domains
        :return:
        """
        _, labels_image = cv2.connectedComponents(self.format_pixel)
        return labels_image

    def _format_pixel(self):
        bg = np.zeros_like(self.text_area_pixel, dtype=np.uint8)
        bg[np.where(self.text_area_pixel > self.pixel_threshold)] = 255
        return bg

    @staticmethod
    def _get_max_size_contours(contours):
        """
        get max siz contours
        """
        return_contours = contours[0]
        length = len(contours)
        for i in range(1, length):
            ret_size = return_contours.shape[0]
            new_size = contours[i].shape[0]
            if ret_size < new_size:
                return_contours = contours[i]
        return return_contours

    def _get_single_bbox(self, i):
        """
        get single bbox
        :param i:
        :return:
        """
        locations = np.where(self.connected_domain == i)
        if locations[0].size < self.min_pixel_size:
            return None
        bg = np.zeros_like(self.connected_domain, dtype=np.uint8)
        bg[locations] = 255
        h, w = bg.shape[0], bg.shape[1]
        tmp_bg = np.zeros((1, w))
        bg = np.concatenate((tmp_bg, bg, tmp_bg), axis=0)
        tmp_bg = np.zeros((h + 2, 1))
        bg = np.concatenate((tmp_bg, bg, tmp_bg), axis=1)
        bg = bg.astype(np.uint8)
        ret, thresh = cv2.threshold(bg, 127, 255, 0)
        contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 0:
            return None
        if len(contours) > 1:
            contours = self._get_max_size_contours(contours)
            contours = [contours]
        rect = cv2.minAreaRect(contours[0])
        box = cv2.boxPoints(rect)
        return box
